from flask import Flask, request, jsonify
from flask_cors import CORS  # 导入 CORS
import pandas as pd
import torch
import torch.nn as nn
import numpy as np
import utils  # 假设 utils 是你自定义的工具模块

# 初始化 Flask 应用
app = Flask(__name__)
CORS(app)  # 启用 CORS 支持

# 定义线性神经网络模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(85, 128)  # 第一个全连接层
        self.dropout = nn.Dropout(p=0.15)  # Dropout层
        self.fc2 = nn.Linear(128, 64)  # 第二个全连接层
        self.fc3 = nn.Linear(64, 2)  # 输出层

    def forward(self, x):
        x = torch.relu(self.fc1(x))  # ReLU激活函数
        x = self.dropout(x)  # Dropout层
        x = torch.relu(self.fc2(x))  # ReLU激活函数
        x = self.dropout(x)  # Dropout层
        x = self.fc3(x)  # 输出层
        return x

# 加载模型
model = Net()
model.load_state_dict(torch.load("./model_epoch_99.pth", map_location=torch.device('cpu')))
model.eval()  # 设置为评估模式

# 定义检测钓鱼网站的API
@app.route('/detect-phishing', methods=['POST'])
def detect_phishing():
    # 获取请求中的URL参数
    data = request.json
    url = data.get('url')
    if not url:
        return jsonify({'error': 'URL is required'}), 400

    try:
        # 使用 utils.getFishX 获取特征数据
        test_data = utils.getFishX(url)
        test_input = test_data.values

        # 检查并转换数据类型
        if test_input.dtype == np.object_:
            test_input = test_input.astype(np.float64)
        else:
            test_input = test_input.astype(np.float64)

        # 转换为 PyTorch 张量
        test_input = torch.from_numpy(test_input).float()

        # 进行预测
        with torch.no_grad():
            output = model(test_input)

        # 解析预测结果
        if output[0].item() > output[1].item():
            result = 0  # 非钓鱼网站
        else:
            result = 1  # 钓鱼网站

        # 返回结果
        return jsonify({'url': url, 'is_phishing': result})

    except Exception as e:
        return jsonify({'error': str(e)}), 500

# 启动Flask应用
if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)